#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <cuda_runtime.h>
#include <png.h>
#include <string>
#include <sstream>
#include <iomanip>

#define WIDTH 1024*4
#define HEIGHT 1024*4
#define RUNS 10

__device__ int julia(int x, int y, int width, int height) {
    const float scale = 1.5f;
    float jx = scale * (float)(width / 2 - x) / (width / 2);
    float jy = scale * (float)(height / 2 - y) / (height / 2);

    float cRe = -0.8f;
    float cIm = 0.156f;

    float zx = jx;
    float zy = jy;

    int max_iter = 200;
    int i;
    for (i = 0; i < max_iter; ++i) {
        float zx2 = zx * zx;
        float zy2 = zy * zy;
        if (zx2 + zy2 > 4.0f) break;
        float temp = zx2 - zy2 + cRe;
        zy = 2.0f * zx * zy + cIm;
        zx = temp;
    }

    return i;
}

__global__ void julia_kernel(unsigned char* ptr, int width, int height) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    if (x >= width || y >= height) return;

    int offset = y * width + x;
    int value = julia(x, y, width, height);
    ptr[offset] = (unsigned char)(255 * value / 200);
}

void write_png(const char* filename, unsigned char* image, int width, int height) {
    FILE* fp = fopen(filename, "wb");
    if (!fp) {
        perror("fopen");
        exit(1);
    }

    png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    png_infop info_ptr = png_create_info_struct(png_ptr);
    if (!png_ptr || !info_ptr) {
        fclose(fp);
        fprintf(stderr, "Failed to create PNG write structs\n");
        exit(1);
    }

    if (setjmp(png_jmpbuf(png_ptr))) {
        fprintf(stderr, "PNG write error\n");
        fclose(fp);
        png_destroy_write_struct(&png_ptr, &info_ptr);
        exit(1);
    }

    png_init_io(png_ptr, fp);
    png_set_IHDR(png_ptr, info_ptr, width, height, 8,
                 PNG_COLOR_TYPE_GRAY, PNG_INTERLACE_NONE,
                 PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
    png_write_info(png_ptr, info_ptr);

    png_bytep* row_pointers = (png_bytep*)malloc(sizeof(png_bytep) * height);
    for (int y = 0; y < height; ++y) {
        row_pointers[y] = image + y * width;
    }

    png_write_image(png_ptr, row_pointers);
    png_write_end(png_ptr, NULL);
    fclose(fp);
    png_destroy_write_struct(&png_ptr, &info_ptr);
    free(row_pointers);
}

int main() {
    int width = WIDTH, height = HEIGHT;
    size_t image_size = width * height;
    float total_time = 0.0f;

    size_t global_peak_mem = 0; // 所有运行的显存峰值

    for (int i = 0; i < RUNS; ++i) {
        printf("\n=== Run %d ===\n", i);

        size_t free_mem_before, total_mem;
        cudaMemGetInfo(&free_mem_before, &total_mem);
        printf("Before Allocation:\n");
        printf("  Free memory:  %.2f MB\n", free_mem_before / (1024.0 * 1024.0));

        unsigned char* d_image;
        unsigned char* h_image = (unsigned char*)malloc(image_size);
        cudaMalloc((void**)&d_image, image_size);

        size_t free_mem_after;
        cudaMemGetInfo(&free_mem_after, &total_mem);
        printf("After Allocation:\n");
        printf("  Free memory:  %.2f MB\n", free_mem_after / (1024.0 * 1024.0));

        size_t used_mem = free_mem_before - free_mem_after;
        printf("Allocated memory: %.2f MB\n", used_mem / (1024.0 * 1024.0));

        // 本轮运行时的峰值显存（初始化为当前占用）
        size_t peak_mem = total_mem - free_mem_after;

        // CUDA Event 用于计时
        cudaEvent_t start, stop;
        cudaEventCreate(&start);
        cudaEventCreate(&stop);

        dim3 block(16, 16);
        dim3 grid((width + block.x - 1) / block.x,
                  (height + block.y - 1) / block.y);

        // 核函数前记录
        size_t mem_before_kernel;
        cudaMemGetInfo(&mem_before_kernel, &total_mem);
        peak_mem = std::max(peak_mem, total_mem - mem_before_kernel);

        cudaEventRecord(start);
        julia_kernel<<<grid, block>>>(d_image, width, height);
        cudaEventRecord(stop);
        cudaEventSynchronize(stop);

        // 核函数后记录
        size_t mem_after_kernel;
        cudaMemGetInfo(&mem_after_kernel, &total_mem);
        peak_mem = std::max(peak_mem, total_mem - mem_after_kernel);

        // 更新全局峰值
        global_peak_mem = std::max(global_peak_mem, peak_mem);

        float ms = 0;
        cudaEventElapsedTime(&ms, start, stop);
        printf("GPU kernel execution time: %.6f ms\n", ms);
        if(i > 0){
            total_time += ms;
        }

        printf("Peak memory usage this run: %.2f MB\n", peak_mem / (1024.0 * 1024.0));

        cudaMemcpy(h_image, d_image, image_size, cudaMemcpyDeviceToHost);

        std::ostringstream filename;
        filename << "img/julia_" << std::setw(2) << std::setfill('0') << i << ".png";
        write_png(filename.str().c_str(), h_image, width, height);
        printf("Saved image to %s\n", filename.str().c_str());

        cudaFree(d_image);
        free(h_image);
        cudaEventDestroy(start);
        cudaEventDestroy(stop);
    }

    printf("\n=== Summary ===\n");
    printf("Average GPU kernel execution time over %d runs: %.6f ms\n", RUNS-1, total_time / (RUNS-1));
    printf("Global peak GPU memory usage over all runs: %.2f MB\n", global_peak_mem / (1024.0 * 1024.0));

    return 0;
}